import json
import random
from faker import Faker

# 初始化Faker库生成随机数据
fake = Faker('zh_CN')

# 定义意图类别
INTENT_CATEGORIES = {
    'CARDIOVASCULAR': ['胸痛', '胸闷', '心悸', '心慌', '心脏不舒服'],
    'TRAUMATIC_INJURY': ['车祸', '摔倒', '摔伤', '撞伤', '割伤'],
    'STROKE': ['说不出话', '面瘫', '肢体麻木', '视力模糊', '头晕'],
    'RESPIRATORY': ['呼吸困难', '哮喘', '喘不过气', '气促', '窒息'],
    'UNCONSCIOUSNESS': ['晕倒', '昏迷', '失去意识', '叫不醒', '晕厥'],
    'PEDIATRIC_FEVER': ['孩子发烧', '婴儿发热', '小儿高热', '宝宝发烧'],
    'POISONING': ['食物中毒', '化学中毒', '药物过量', '误食'],
    'BURN': ['烫伤', '烧伤', '着火', '火灾'],
    'FRACTURE': ['骨折', '骨头断了', '肢体变形', '不能动'],
    'DROWNING': ['溺水', '落水', '淹水', '游泳事故']
}

# 定义症状词汇
SYMPTOMS = {
    'SYMPTOM': [
        '疼痛', '剧痛', '刺痛', '绞痛', '呼吸困难', '气促', '喘息', '咳嗽',
        '咳血', '呕吐', '恶心', '腹泻', '便血', '头晕', '头痛', '眩晕',
        '昏迷', '晕厥', '抽搐', '痉挛', '麻木', '无力', '瘫痪', '视力模糊',
        '失明', '听力下降', '耳鸣', '心悸', '心慌', '胸闷', '胸痛', '腹痛',
        '出血', '流血', '吐血', '便血', '血尿', '发热', '发烧', '寒战',
        '出汗', '盗汗', '冷汗', '皮疹', '瘙痒', '红肿', '肿胀', '变形',
        '骨折', '脱臼', '烧伤', '烫伤', '伤口', '撕裂伤', '挫伤'
    ],
    'LOCATION': [
        '家里', '办公室', '学校', '商场', '公园', '街道', '马路', '高速公路',
        '餐厅', '厨房', '卫生间', '卧室', '客厅', '小区', '地铁站', '公交站',
        '火车站', '机场', '游泳池', '海滩', '山上', '野外'
    ],
    'DURATION': [
        '几分钟', '十分钟', '半小时', '一小时', '两小时', '三小时', '半天',
        '一天', '两天', '三天', '一周', '刚才', '刚刚', '现在'
    ],
    'COUNT': [
        '一个人', '两个人', '三个人', '多人', '几个', '很多', '大量'
    ]
}


def generate_emergency_data(num_samples=1000):
    data = []

    for i in range(num_samples):
        # 随机选择意图类别
        intent = random.choice(list(INTENT_CATEGORIES.keys()))

        # 根据意图生成相关文本
        if intent == 'CARDIOVASCULAR':
            text = f"我{fake.random_element(['父亲', '母亲', '爷爷', '奶奶'])}突然{fake.random_element(INTENT_CATEGORIES[intent])}，{fake.random_element(['冒冷汗', '呼吸困难', '脸色苍白'])}，持续{fake.random_element(SYMPTOMS['DURATION'])}了"
            entities = [
                {'word': fake.random_element(INTENT_CATEGORIES[intent]), 'label': 'SYMPTOM'},
                {'word': fake.random_element(['冒冷汗', '呼吸困难', '脸色苍白']), 'label': 'SYMPTOM'},
                {'word': fake.random_element(SYMPTOMS['DURATION']), 'label': 'DURATION'}
            ]

        elif intent == 'TRAUMATIC_INJURY':
            location = fake.random_element(SYMPTOMS['LOCATION'])
            text = f"在{location}发生{fake.random_element(['车祸', '事故', '意外'])}, {fake.random_element(SYMPTOMS['COUNT'])}受伤，{fake.random_element(['流血', '骨折', '昏迷'])}"
            entities = [
                {'word': location, 'label': 'LOCATION'},
                {'word': fake.random_element(SYMPTOMS['COUNT']), 'label': 'COUNT'},
                {'word': fake.random_element(['流血', '骨折', '昏迷']), 'label': 'SYMPTOM'}
            ]

        elif intent == 'STROKE':
            text = f"{fake.random_element(['我父亲', '我母亲', '患者'])}突然{fake.random_element(INTENT_CATEGORIES[intent])}，{fake.random_element(['说话不清', '嘴角歪斜', '肢体无力'])}"
            entities = [
                {'word': fake.random_element(INTENT_CATEGORIES[intent]), 'label': 'SYMPTOM'},
                {'word': fake.random_element(['说话不清', '嘴角歪斜', '肢体无力']), 'label': 'SYMPTOM'}
            ]

        elif intent == 'RESPIRATORY':
            text = f"{fake.random_element(['哮喘发作', '呼吸困难', '喘不过气'])}，{fake.random_element(['嘴唇发紫', '需要氧气', '无法平躺'])}"
            entities = [
                {'word': fake.random_element(['哮喘发作', '呼吸困难', '喘不过气']), 'label': 'SYMPTOM'},
                {'word': fake.random_element(['嘴唇发紫', '需要氧气', '无法平躺']), 'label': 'SYMPTOM'}
            ]

        else:
            # 为其他意图生成通用文本
            symptom = fake.random_element(SYMPTOMS['SYMPTOM'])
            duration = fake.random_element(SYMPTOMS['DURATION'])
            text = f"{fake.random_element(['突然', '刚刚', '现在'])}{symptom}，已经持续{duration}，{fake.random_element(['很难受', '很严重', '需要帮助'])}"
            entities = [
                {'word': symptom, 'label': 'SYMPTOM'},
                {'word': duration, 'label': 'DURATION'}
            ]

        # 随机添加地点信息
        if random.random() > 0.3:  # 70%的概率添加地点
            location = fake.random_element(SYMPTOMS['LOCATION'])
            text = f"在{location}，" + text
            entities.append({'word': location, 'label': 'LOCATION'})

        data.append({
            'text': text,
            'entities': entities,
            'intent': intent
        })

    return data


# 生成1000条数据
emergency_data = generate_emergency_data(10000)

# 保存到JSON文件
with open('../datasets/emergency_calls_1000.json', 'w', encoding='utf-8') as f:
    json.dump(emergency_data, f, ensure_ascii=False, indent=2)

print(f"已生成1000条急救标注数据，保存至 emergency_calls_1000.json")